use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::TypeEnvironment;
use super::super::TypeInference;

impl TypeInference {



    /// Infer types in an array expression
    pub(in crate::type_check::type_inference) fn infer_array_expression(
        &mut self,
        array: &ArrayExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        if array.elements.is_empty() {
            return Err(CompilerError::type_error(
                array.location.line,
                array.location.column,
                "Cannot infer type of empty array",
            ));
        }

        // Infer type of first element
        let first_type = self.infer_expression(&array.elements[0], env)?;

        // Check that all elements have the same type
        for element in &array.elements[1..] {
            let element_type = self.infer_expression(element, env)?;
            self.add_constraint(&element_type, &first_type, array.location)?;
        }

        Ok(Type::Array(ArrayType {
            element_type: Box::new(first_type),
            size: None,
            location: array.location,
        }))
    }



    /// Infer types in a map expression
    pub(in crate::type_check::type_inference) fn infer_map_expression(&mut self, map: &MapExpr, env: &mut TypeEnvironment) -> Result<Type> {
        if map.entries.is_empty() {
            return Err(CompilerError::type_error(
                map.location.line,
                map.location.column,
                "Cannot infer type of empty map",
            ));
        }

        // Infer types of first entry
        let first_key_type = self.infer_expression(&map.entries[0].key, env)?;
        let first_value_type = self.infer_expression(&map.entries[0].value, env)?;

        // Check that all entries have the same types
        for entry in &map.entries[1..] {
            let key_type = self.infer_expression(&entry.key, env)?;
            let value_type = self.infer_expression(&entry.value, env)?;
            self.add_constraint(&key_type, &first_key_type, map.location)?;
            self.add_constraint(&value_type, &first_value_type, map.location)?;
        }

        Ok(Type::Map(MapType {
            key_type: Box::new(first_key_type),
            value_type: Box::new(first_value_type),
            location: map.location,
        }))
    }



    /// Infer types in a struct expression
    pub(in crate::type_check::type_inference) fn infer_struct_expression(
        &mut self,
        struct_expr: &StructExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        // Infer field types
        for field in &struct_expr.fields {
            self.infer_expression(&field.value, env)?;
        }

        Ok(struct_expr.struct_type.clone())
    }

}
